import pygame, sys
import numpy as np
import random
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
BLUE = (0, 100, 255)
RED = (255,0,0)
GREY = (230, 230, 230)
GREEN=(0,255,0)
BACKGROUND = BLACK
class Dot(pygame.sprite.Sprite):
def __init__(
self,
x,
y,
width,
height,
color=BLACK,
radius=5,
velocity=[0, 0],
randomize=False,
):
super().__init__()
self.image = pygame.Surface([radius * 2, radius * 2])
self.image.fill(BACKGROUND)
pygame.draw.circle(
self.image, color, (radius, radius), radius
)
self.rect = self.image.get_rect()
self.pos = np.array([x, y], dtype=np.float64)
self.vel = np.asarray(velocity, dtype=np.float64)
self.killswitch_on = False
self.recovered = False
self.randomize = randomize
self.WIDTH = width
self.HEIGHT = height
def update(self): # se un dot esce dal riquadro rientra dalla parte opposta
self.pos += self.vel
x, y = self.pos
if x < 0:
self.pos[0] = self.WIDTH
x = self.WIDTH
if x > self.WIDTH:
self.pos[0] = 0
x = 0
if y < 0:
self.pos[1] = self.HEIGHT
y = self.HEIGHT
if y > self.HEIGHT:
self.pos[1] = 0
y = 0
self.rect.x = x
self.rect.y = y
vel_norm = np.linalg.norm(self.vel)
if vel_norm > 3:
self.vel /= vel_norm
if self.randomize:
self.vel += np.random.rand(2) * 2 - 1 #-1/1 gfg
if self.killswitch_on:
self.cycles -= 1
if self.cycles<= 0:
self.killswitch_on = False # non si può morire 2 volte
some_number = np.random.rand()
if self.mortality_rate > some_number:
self.kill()
else:
self.recovered = True
def respawn(self, color, radius=5):
return Dot(
self.rect.x,
self.rect.y,
self.WIDTH,
self.HEIGHT,
color=color,
velocity=self.vel,
)
def killswitch(self, cycles=20, mortality_rate=0.2): # set di update di ciclo
self.killswitch_on = True
self.cycles = cycles
self.mortality_rate = mortality_rate
class Explosion(pygame.sprite.Sprite):
def __init__(self, x,y ):
pygame.sprite.Sprite.__init__(self)
self.x = x
self.y = y
self.images = []
for num in range(0,9):
img=pygame.image.load(f"regularExplosion0{num}.png")
img= pygame.transform.scale(img, (100,100))
self.images.append(img)
self.index = 0
self.image=self .images[self.index]
self.rect = self.image.get_rect()
self.rect.center = [x,y]
self.counter = 0
def update(self):
explosions_speed=4
self.counter +=1
if self.counter>=explosions_speed and self.index < len(self.images)-1:
self.counter=0
self.index+=1
self.image=self.images[self.index]
if self.index>= len(self.images) -1 and self.counter>=explosions_speed:
self.kill()
explosions_group = pygame.sprite.Group()
class Simulation:
def __init__(self, width=600, height=480):
self.WIDTH = width
self.HEIGHT = height
self.susceptible_container = pygame.sprite.Group()
self.infected_container = pygame.sprite.Group()
self.recovered_container = pygame.sprite.Group()
self.all_container = pygame.sprite.Group()
self.n_susceptible = 20
self.n_infected = 1
self.n_died=0
self.n_infected = 0
self.T = 1000
self.cycles = 20
self.mortality_rate = 0.2
def start(self, randomize=False):
self.N = (
self.n_susceptible + self.n_infected + self.n_quarantined
)
pygame.init()
screen = pygame.display.set_mode([self.WIDTH, self.HEIGHT])
for i in range(self.n_susceptible):
x = np.random.randint(0, self.WIDTH + 1)
y = np.random.randint(0, self.HEIGHT + 1)
vel = np.random.rand(2) * 2 - 1
guy = Dot(
x,
y,
self.WIDTH,
self.HEIGHT,
color=WHITE,
velocity=vel,
randomize=randomize,
)
self.susceptible_container.add(guy)
self.all_container.add(guy)
for i in range(self.n_quarantined):
x = np.random.randint(0, self.WIDTH + 1)
y = np.random.randint(0, self.HEIGHT + 1)
vel = [0, 0]
guy = Dot(
x,
y,
self.WIDTH,
self.HEIGHT,
color=BLUE,
velocity=vel,
randomize=False,
)
self.susceptible_container.add(guy)
self.all_container.add(guy)
for i in range(self.n_infected):
x = np.random.randint(0, self.WIDTH + 1)
y = np.random.randint(0, self.HEIGHT + 1)
vel = np.random.rand(2) * 2 - 1
guy = Dot(
x,
y,
self.WIDTH,
self.HEIGHT,
color=RED,
velocity=vel,
randomize=randomize,
)
self.infected_container.add(guy)
self.all_container.add(guy)
clock = pygame.time.Clock()
for i in range(self.T):
for event in pygame.event.get():
if event.type == pygame.QUIT:
sys.exit()
self.all_container.update()
screen.fill(BACKGROUND)
# Infetti
collision_group = pygame.sprite.groupcollide( # controllo collisione
self.susceptible_container,
self.infected_container,
True, #remove susc
False, #lascia inf
)
for guy in collision_group:
new_guy = guy.respawn(RED) #dobbiamo cambiare il colore -> rimuovo da S aggiungo a I
new_guy.vel *= -1
new_guy.killswitch(
self.cycles, self.mortality_rate
)
self.infected_container.add(new_guy)
self.all_container.add(new_guy)
# rimessi
recovered = []
for guy in self.infected_container:
some_number = np.random.rand()
if guy.recovered and self.mortality_rate > some_number:
new_guy = Explosion(random.randrange(0, 600),random.randrange(0, 480))
self.recovered_container.add(new_guy)
self.all_container.add(new_guy)
recovered.append(guy)
for guy in self.infected_container:
if guy.recovered:
new_guy = guy.respawn(GREEN)
self.recovered_container.add(new_guy)
self.all_container.add(new_guy)
recovered.append(guy)
if len(recovered) > 0:
self.infected_container.remove(*recovered) # leviamo i recovered da infetti
self.all_container.remove(*recovered)
self.all_container.draw(screen)
pygame.display.flip()
clock.tick(30)
pygame.quit()
if __name__ == "__main__":
scelto = int(input('premere:\n"0" per inserire manualmente i parametri della malattia\n"1" per usare dei paramteri preimpostati \nScelto: '))
if scelto == 0:
SIR = Simulation(600, 480)
SIR.n_susceptible = int(input('inserire il numero di persone sane: '))
SIR.n_quarantined = 0
SIR.n_infected = int(input('inserire il numero di persone infette: '))
SIR.cycles= 200 # indica sopo quanto tempo i dot possono iniziare a passare a R
SIR.mortality_rate = float(input('inserire il tasso di mortalità: '))
SIR.start(randomize=True)
elif scelto==1:
SIR = Simulation(600, 480)
SIR.n_susceptible = 100
SIR.n_quarantined = 0
SIR.n_infected = 5
SIR.cycles= 200
SIR.mortality_rate = 0.2
SIR.start(randomize=True)
else:
print('errore')
